'''
MDAUNet: Multi-Domain Aware Unet: For multi-domain incremental change detection
Can dynamicly increase the model
'''

import torch.nn as nn
import torch

current_task = 0
class conv_block_nested_DA(nn.Module):
    def __init__(self, in_ch, mid_ch, out_ch, nb_tasks):
        super(conv_block_nested_DA, self).__init__()
        self.activation = nn.ReLU(inplace=True)
        self.conv1 = nn.Conv2d(in_ch, mid_ch, kernel_size=3, padding=1, bias=True)
        self.bn1 = nn.BatchNorm2d(mid_ch)
        self.conv2 = nn.Conv2d(mid_ch, out_ch, kernel_size=3, padding=1, bias=True)
        self.bn2 = nn.BatchNorm2d(out_ch)

        self.parallel_conv_1 = nn.ModuleList(
            [nn.Conv2d(in_ch, mid_ch, kernel_size=1, stride=1, padding=0, bias=True) for i in range(nb_tasks)])
        self.bns_1 = nn.ModuleList([nn.BatchNorm2d(mid_ch, eps=1e-03) for i in range(nb_tasks)])
        self.parallel_conv_2 = nn.ModuleList(
            [nn.Conv2d(mid_ch, out_ch, kernel_size=1, stride=1, padding=0, bias=True) for i in range(nb_tasks)])
        self.bns_2 = nn.ModuleList([nn.BatchNorm2d(out_ch, eps=1e-03) for i in range(nb_tasks)])


    def forward(self, input):
        task = current_task
        # print("1", x.shape)
        output = self.conv1(input)  # share conv1

        output = output + self.parallel_conv_1[task](input)
        identity = output

        output = self.bns_1[task](output)
        output_ = self.activation(output)


        output = self.conv2(output_)
        output = output + self.parallel_conv_2[task](output_)

        output = self.bns_2[task](output)
        output = self.activation(output + identity)
        # print("out.shape", output.shape)
        return output

class conv_block_nested(nn.Module):
    def __init__(self, in_ch, mid_ch, out_ch):
        super(conv_block_nested, self).__init__()
        self.activation = nn.ReLU(inplace=True)
        self.conv1 = nn.Conv2d(in_ch, mid_ch, kernel_size=3, padding=1, bias=True)
        self.bn1 = nn.BatchNorm2d(mid_ch)
        self.conv2 = nn.Conv2d(mid_ch, out_ch, kernel_size=3, padding=1, bias=True)
        self.bn2 = nn.BatchNorm2d(out_ch)

    def forward(self, x):
        # print("1", x.shape)
        x = self.conv1(x)
        # print("2", x.shape)
        identity = x
        x = self.bn1(x)
        x = self.activation(x)

        # print("3", x.shape)
        x = self.conv2(x)
        # print("4", x.shape)
        x = self.bn2(x)
        output = self.activation(x + identity)
        # print("out.shape", output.shape)
        return output

class up(nn.Module):
    def __init__(self, in_ch, bilinear=False):
        super(up, self).__init__()

        if bilinear:
            self.up = nn.Upsample(scale_factor=2,
                                  mode='bilinear',
                                  align_corners=True)
        else:
            self.up = nn.ConvTranspose2d(in_ch, in_ch, 2, stride=2)

    def forward(self, x):
        x = self.up(x)
        return x

class ChannelAttention(nn.Module):
    def __init__(self, in_channels, ratio = 16):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.fc1 = nn.Conv2d(in_channels,in_channels//ratio,1,bias=False)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Conv2d(in_channels//ratio, in_channels,1,bias=False)
        self.sigmod = nn.Sigmoid()
    def forward(self,x):
        avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x))))
        max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))
        out = avg_out + max_out
        return self.sigmod(out)


n1 = 32
filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]
class Encoder(nn.Module):
    def __init__(self, nb_tasks=1):
        super().__init__()
        in_ch = 3
        self.conv0_0 = conv_block_nested_DA(in_ch, filters[0], filters[0], nb_tasks)
        self.conv1_0 = conv_block_nested_DA(filters[0], filters[1], filters[1], nb_tasks)
        self.conv2_0 = conv_block_nested_DA(filters[1], filters[2], filters[2], nb_tasks)
        self.conv3_0 = conv_block_nested_DA(filters[2], filters[3], filters[3], nb_tasks)
        self.conv4_0 = conv_block_nested_DA(filters[3], filters[4], filters[4], nb_tasks)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

    def forward(self, input):
        x0_0 = self.conv0_0(input)
        x1_0 = self.conv1_0(self.pool(x0_0))
        x2_0 = self.conv2_0(self.pool(x1_0))
        x3_0 = self.conv3_0(self.pool(x2_0))
        x4_0 = self.conv4_0(self.pool(x3_0))
        return [x0_0, x1_0, x2_0, x3_0, x4_0]


class Decoder(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        out_ch = 2
        self.conv3_1 = conv_block_nested(filters[3] * 2 + filters[4], filters[3], filters[3])
        self.conv2_2 = conv_block_nested(filters[2] * 2 + filters[3], filters[2], filters[2])
        self.conv1_3 = conv_block_nested(filters[1] * 2 + filters[2], filters[1], filters[1])
        self.conv0_4 = conv_block_nested(filters[0] * 2 + filters[1], filters[0], filters[0])
        self.Up4_0 = up(filters[4])
        self.Up3_1 = up(filters[3])
        self.Up2_2 = up(filters[2])
        self.Up1_3 = up(filters[1])
        self.conv_final = nn.Conv2d(filters[0], out_ch, kernel_size=1)

    def forward(self, x_A, x_B):
        self.x0_0A, self.x1_0A, self.x2_0A, self.x3_0A, self.x4_0A = x_A
        self.x0_0B, self.x1_0B, self.x2_0B, self.x3_0B, self.x4_0B = x_B

        x3_1 = self.conv3_1(torch.cat([self.x3_0A, self.x3_0B, self.Up4_0(self.x4_0B)], 1))
        x2_2 = self.conv2_2(torch.cat([self.x2_0A, self.x2_0B, self.Up3_1(x3_1)], 1))
        x1_3 = self.conv1_3(torch.cat([self.x1_0A, self.x1_0B, self.Up2_2(x2_2)], 1))
        x0_4 = self.conv0_4(torch.cat([self.x0_0A, self.x0_0B, self.Up1_3(x1_3)], 1))
        out = self.conv_final(x0_4)
        return out


class MDAUNet(nn.Module):
    def __init__(self, num_classes=[2], nb_tasks=1, cur_task=0):
        super(MDAUNet, self).__init__()
        torch.nn.Module.dump_patches = True
        global current_task
        current_task = cur_task
        print('hi, inside erfnet_RA_parallel', current_task, nb_tasks)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

        '''Encoder return multi feature map to nested
        '''
        self.encoder = Encoder(nb_tasks)
        self.decoder = nn.ModuleList([Decoder(num_classes[i]) for i in range(nb_tasks)])


    def forward(self, input, task):
        global current_task
        current_task = task

        xA, xB = input[:,:3,:,:], input[:,3:,:,:]
        x0_A = self.encoder(xA)
        x0_B = self.encoder(xB)
        out = self.decoder[task].forward(x0_A, x0_B)
        return out

if __name__ == "__main__":
    num_classes = [2,2]
    current_task = 1
    unet = MDAUNet(nb_tasks=2, cur_task=current_task, num_classes=num_classes)
    image_A = torch.randn((4, 3, 256, 256))
    image_B = torch.randn((4, 3, 256, 256))
    for name, m in unet.named_parameters():
        print(name)
    outputs_change, feature_map = unet(image_A, image_B, current_task)
